Skip to content

feat(rust): Laguna-S-2.1 decode on the CUDA engine, 36.5 tok/s on GB10#69

Merged
TheTom merged 2 commits into
tom/feat/cuda-hip-vulkan-backendsfrom
tom/laguna-rust-engine
Jul 23, 2026
Merged

feat(rust): Laguna-S-2.1 decode on the CUDA engine, 36.5 tok/s on GB10#69
TheTom merged 2 commits into
tom/feat/cuda-hip-vulkan-backendsfrom
tom/laguna-rust-engine

Conversation

@TheTom

@TheTom TheTom commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

What

Full single-stream Laguna-S-2.1 (117.55B / 8.14B-active hybrid MoE) decode on the Rust CUDA engine, correctness-verified token-for-token against the reference C++ implementation and benched on the DGX Spark GB10.

Commit 1 lands the in-flight engine state the GB10 working tree has been running (CUDA-graph trait plumbing, fp8 microbench, MoE grouped-MMA test updates). Commit 2 is the Laguna port itself.

Numbers (GB10, single-stream greedy, 128-token bench, Q4)

stack tok/s
reference C++ implementation (same weights, same box) 21.24
model vendor's stated no-spec ceiling, any engine 13-14
this PR, eager 28-30
this PR, graph replay 34.1
this PR, graph + fused MoE gather 35.0
this PR, graph + fuse + micro-fusions 36.5

Greedy continuations byte-identical to the reference in every mode. Roughly 62 percent of the theoretical bandwidth ceiling for 8.14B active params at Q4 on this memory system.

How

  • Q4 requant load straight from GGUF (Q8 LM head), YaRN scaling read from GGUF metadata (this export is a 256K / factor-32 checkpoint, not the model card's 1M / 128)
  • Sigmoid top-10 router with score-correction bias, device-side gather-GEMV expert path (indices resolved in-kernel), shared expert, per-head softplus attention gate, per-head QK RMSNorm, period-4 hybrid attention with a 512-slot ring KV cache
  • Whole-step CUDA graph capture (warmup-capture-replay, persistent u32 workspace, n_kv and rope position read from device buffers)
  • Fused MoE gate+up+SwiGLU gather; concatenated QKVG projection; shared expert through the fused gather; o_proj residual accumulate
  • Env-gated: FFAI_LAGUNA_GGUF (weights), FFAI_LAGUNA_GRAPH / _FUSE / _MICRO / _FAST / _GEN, debug dumps via _DEBUG

Tests

  • Kernel unit tests run without weights (rope yarn pos-0 invariants at model geometry, softplus gate, nbuf sdpa vs registry kernel at gqa 6 and 9, fused swiglu gather vs unfused compose)
  • Env-gated integration smoke: greedy generation with detokenized output and tok/s
  • Warp-count sweep harness (FFAI_LAGUNA_SWEEP) benches launch-config combos on one model load

Notes for review

  • Companion kernel-side PR in the metaltile tree adds ffai_moe_gather_q4 / _swiglu and the bf16 casts this branch calls
  • The sync commit is working-tree state predating this feature; review separately
  • Known cosmetic debt: debug instrumentation in laguna.rs is env-gated but verbose; happy to strip on request

TheTom added 2 commits July 22, 2026 22:42
Carries the uncommitted engine state the GB10 box has been running:
CUDA-graph capture trait plumbing (begin/end capture, graph_launch),
fp8 projection microbench and MoE grouped-MMA test updates, and the
device trait additions they depend on. Precedes the Laguna port commits
that build on these interfaces.
Full single-stream decode for the 117.55B/8.14B-active hybrid MoE:
sigmoid top-10 router with score-correction bias and shared expert,
per-head softplus attention gate, per-head QK RMSNorm, period-4
full/sliding-window attention with a 512-slot ring KV cache, YaRN
partial-rotary rope on full layers (scaling read from GGUF metadata,
this export is a 256K/factor-32 checkpoint), Q8 LM head, and a
Q4-requantized weight load straight from GGUF.

New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead,
argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer),
write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers.

Decode runs eagerly or as a captured CUDA graph (FFAI_LAGUNA_GRAPH=1,
warmup-capture-replay via LagunaDecodeCtx with a persistent workspace),
with optional fused MoE gate+up (FFAI_LAGUNA_FUSE=1) and micro-fusions
(FFAI_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via
the fused gather, o_proj residual accumulate).

GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s
here (graph+fuse+micro), greedy tokens verified identical to the reference
implementation. Env-gated integration smoke: FFAI_LAGUNA_GGUF plus
laguna tests in ffai-cuda (unit kernels run without weights).
@github-actions github-actions Bot added the feature New feature or capability label Jul 23, 2026
@TheTom
TheTom marked this pull request as ready for review July 23, 2026 12:10
@TheTom
TheTom merged commit 2e951ef into tom/feat/cuda-hip-vulkan-backends Jul 23, 2026
2 checks passed
TheTom added a commit that referenced this pull request Jul 23, 2026
#69) (renamed)

* sync(rust): land in-flight CUDA engine work from the GB10 working tree

Carries the uncommitted engine state the GB10 box has been running:
CUDA-graph capture trait plumbing (begin/end capture, graph_launch),
fp8 projection microbench and MoE grouped-MMA test updates, and the
device trait additions they depend on. Precedes the Laguna port commits
that build on these interfaces.

* feat(laguna): Laguna-S-2.1 decode on the CUDA engine

Full single-stream decode for the 117.55B/8.14B-active hybrid MoE:
sigmoid top-10 router with score-correction bias and shared expert,
per-head softplus attention gate, per-head QK RMSNorm, period-4
full/sliding-window attention with a 512-slot ring KV cache, YaRN
partial-rotary rope on full layers (scaling read from GGUF metadata,
this export is a 256K/factor-32 checkpoint), Q8 LM head, and a
Q4-requantized weight load straight from GGUF.

New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead,
argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer),
write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers (kept as
always-error stubs in this rebrand — the iron_moe_gather_q4 /
iron_moe_gather_q4_swiglu kernels aren't in thewafflehaus/iron@dev yet).

Decode runs eagerly or as a captured CUDA graph (BUTTER_LAGUNA_GRAPH=1,
warmup-capture-replay via LagunaDecodeCtx with a persistent workspace),
with optional fused MoE gate+up (BUTTER_LAGUNA_FUSE=1) and micro-fusions
(BUTTER_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via
the fused gather, o_proj residual accumulate).

GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s
here (graph+fuse+micro), greedy tokens verified identical to the reference
implementation. Env-gated integration smoke: BUTTER_LAGUNA_GGUF plus
laguna tests in wh-butter-cuda (unit kernels run without weights).
TheTom added a commit that referenced this pull request Jul 23, 2026
…gates

Both cherry-picked commits (#69, #70) brought every op they need with
them — no separate ops-porting commit was required. These are gate-driven
fixups only:

- wh-butter-cuda/src/imp.rs: drop the moe_marlin_gemm/marlin_repack/
  marlin_build_routing pass-throughs PR #70 added — they call through to
  methods on wh_iron_runtime::CudaDevice that don't exist in
  thewafflehaus/iron@dev (confirmed: no `marlin` symbol anywhere in that
  repo). Falls back to the wh-butter-core default "unsupported on this
  backend" stub (already present, unaffected) instead of failing to
  compile. Same root cause as the moe_gather_q4/moe_gather_q4_swiglu stubs
  already documented in wh-butter-ops.
- wh-butter-cuda/tests/all_models.rs: dropped a `smallmodel_fuse_slicecast_ab`
  test calling `wh_butter_modeltests::bench_smallmodel_fuse_slicecast` —
  confirmed that function doesn't exist anywhere in wh-butter-modeltests
  on the original tom/feat/cuda-hip-vulkan-backends branch either
  (pre-existing dead reference predating this cherry-pick, not something
  it introduced).
- wh-butter-cuda/tests/f16norm_f32in.rs: dropped — same pre-existing-bug
  class, references a `add_rms_norm_f16norm` op that was never
  implemented on the source branch.
- rust/Cargo.lock: regenerated via `cargo build` (not hand-edited).
TheTom added a commit that referenced this pull request Jul 24, 2026
#69) (renamed)

* sync(rust): land in-flight CUDA engine work from the GB10 working tree

Carries the uncommitted engine state the GB10 box has been running:
CUDA-graph capture trait plumbing (begin/end capture, graph_launch),
fp8 projection microbench and MoE grouped-MMA test updates, and the
device trait additions they depend on. Precedes the Laguna port commits
that build on these interfaces.

* feat(laguna): Laguna-S-2.1 decode on the CUDA engine

Full single-stream decode for the 117.55B/8.14B-active hybrid MoE:
sigmoid top-10 router with score-correction bias and shared expert,
per-head softplus attention gate, per-head QK RMSNorm, period-4
full/sliding-window attention with a 512-slot ring KV cache, YaRN
partial-rotary rope on full layers (scaling read from GGUF metadata,
this export is a 256K/factor-32 checkpoint), Q8 LM head, and a
Q4-requantized weight load straight from GGUF.

New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead,
argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer),
write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers (kept as
always-error stubs in this rebrand — the iron_moe_gather_q4 /
iron_moe_gather_q4_swiglu kernels aren't in thewafflehaus/iron@dev yet).

Decode runs eagerly or as a captured CUDA graph (BUTTER_LAGUNA_GRAPH=1,
warmup-capture-replay via LagunaDecodeCtx with a persistent workspace),
with optional fused MoE gate+up (BUTTER_LAGUNA_FUSE=1) and micro-fusions
(BUTTER_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via
the fused gather, o_proj residual accumulate).

GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s
here (graph+fuse+micro), greedy tokens verified identical to the reference
implementation. Env-gated integration smoke: BUTTER_LAGUNA_GGUF plus
laguna tests in wh-butter-cuda (unit kernels run without weights).
TheTom added a commit that referenced this pull request Jul 24, 2026
…gates

Both cherry-picked commits (#69, #70) brought every op they need with
them — no separate ops-porting commit was required. These are gate-driven
fixups only:

- wh-butter-cuda/src/imp.rs: drop the moe_marlin_gemm/marlin_repack/
  marlin_build_routing pass-throughs PR #70 added — they call through to
  methods on wh_iron_runtime::CudaDevice that don't exist in
  thewafflehaus/iron@dev (confirmed: no `marlin` symbol anywhere in that
  repo). Falls back to the wh-butter-core default "unsupported on this
  backend" stub (already present, unaffected) instead of failing to
  compile. Same root cause as the moe_gather_q4/moe_gather_q4_swiglu stubs
  already documented in wh-butter-ops.
- wh-butter-cuda/tests/all_models.rs: dropped a `smallmodel_fuse_slicecast_ab`
  test calling `wh_butter_modeltests::bench_smallmodel_fuse_slicecast` —
  confirmed that function doesn't exist anywhere in wh-butter-modeltests
  on the original tom/feat/cuda-hip-vulkan-backends branch either
  (pre-existing dead reference predating this cherry-pick, not something
  it introduced).
- wh-butter-cuda/tests/f16norm_f32in.rs: dropped — same pre-existing-bug
  class, references a `add_rms_norm_f16norm` op that was never
  implemented on the source branch.
- rust/Cargo.lock: regenerated via `cargo build` (not hand-edited).
TheTom added a commit that referenced this pull request Jul 24, 2026
…#70) (#73)

* feat(rust): Laguna-S-2.1 decode on the CUDA engine, 36.5 tok/s on GB10 (#69) (renamed)

* sync(rust): land in-flight CUDA engine work from the GB10 working tree

Carries the uncommitted engine state the GB10 box has been running:
CUDA-graph capture trait plumbing (begin/end capture, graph_launch),
fp8 projection microbench and MoE grouped-MMA test updates, and the
device trait additions they depend on. Precedes the Laguna port commits
that build on these interfaces.

* feat(laguna): Laguna-S-2.1 decode on the CUDA engine

Full single-stream decode for the 117.55B/8.14B-active hybrid MoE:
sigmoid top-10 router with score-correction bias and shared expert,
per-head softplus attention gate, per-head QK RMSNorm, period-4
full/sliding-window attention with a 512-slot ring KV cache, YaRN
partial-rotary rope on full layers (scaling read from GGUF metadata,
this export is a 256K/factor-32 checkpoint), Q8 LM head, and a
Q4-requantized weight load straight from GGUF.

New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead,
argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer),
write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers (kept as
always-error stubs in this rebrand — the iron_moe_gather_q4 /
iron_moe_gather_q4_swiglu kernels aren't in thewafflehaus/iron@dev yet).

Decode runs eagerly or as a captured CUDA graph (BUTTER_LAGUNA_GRAPH=1,
warmup-capture-replay via LagunaDecodeCtx with a persistent workspace),
with optional fused MoE gate+up (BUTTER_LAGUNA_FUSE=1) and micro-fusions
(BUTTER_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via
the fused gather, o_proj residual accumulate).

GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s
here (graph+fuse+micro), greedy tokens verified identical to the reference
implementation. Env-gated integration smoke: BUTTER_LAGUNA_GGUF plus
laguna tests in wh-butter-cuda (unit kernels run without weights).

* feat(rust): Laguna batched prefill, 156 to 559 pp tok/s, plus weight cache and dtype-cache fix (#70) (renamed)

* fix(cuda): key compiled-module and shared-size caches by kernel dtype signature

The ops layer caches kernel IR by (name, dtype) but the backend cached
compiled modules by bare kernel name, so the first dtype to touch a name
won and every other dtype silently ran the wrong binary: wrong element
stride, out-of-bounds reads, and in the shrinking-stride direction silent
corruption. Found when the prefill path's first f16 gather inherited the
decode path's f32 module. Shared-memory sizing had the same hazard.

* feat(laguna): batched prefill with tensor-core projections, fused MoE, CTA scheduling

Chunked multi-token prefill (default chunk 2048): batched YaRN/plain rope,
per-query windowed varlen attention with a window-aware KV-block skip,
linear sliding-window scratch compacted into the decode ring, grouped-GEMM
MoE with on-device descriptors, Marlin W4A16 tensor-core dense projections
(concatenated QKV, o_proj, dense FFN, shared expert), fused gate+up expert
stacks, and grouped-GEMM CTA scheduling (descending-size expert order plus
N-banded CTA order for weight L2 reuse, default on).

Correctness gates: prefill-then-decode greedy continuation byte-identical
to decode-only; last-token argmax matches the reference oracle; kernel
unit tests for the batched rope, batched gate, windowed varlen skip, fused
swiglu gather, and scheduling A/B on skewed synthetic groups.

GB10 single-stream: prefill 559 tok/s at 2048 (was 156 at first light),
504 at 8192; decode unchanged at 36.3 via graph replay. The reference C++ engine
on identical weights and box: 663 and 660 stock.

* feat(laguna): on-disk weight cache and windowed parallel conversion

Content-keyed cache of every converted engine-format weight blob (per
tensor artifact, keyed by format version, source GGUF identity, and the
load-shaping env flags), written on first conversion and mmap-read on
later loads. Conversion itself runs rayon-parallel over a bounded window
of layers (full parallelism held tens of GB of transients and got
OOM-killed on the shared 128GB). Warm reload: 56s, down from ~7.5
minutes; cache hits and misses are reported at load end.

Also carries the comment hygiene sweep across the Laguna files (dash
style, neutral phrasing for external references) and the stale
decode-only module doc fix.

* fix(rust): iron-dependency gaps + pre-existing test bugs surfaced by gates

Both cherry-picked commits (#69, #70) brought every op they need with
them — no separate ops-porting commit was required. These are gate-driven
fixups only:

- wh-butter-cuda/src/imp.rs: drop the moe_marlin_gemm/marlin_repack/
  marlin_build_routing pass-throughs PR #70 added — they call through to
  methods on wh_iron_runtime::CudaDevice that don't exist in
  thewafflehaus/iron@dev (confirmed: no `marlin` symbol anywhere in that
  repo). Falls back to the wh-butter-core default "unsupported on this
  backend" stub (already present, unaffected) instead of failing to
  compile. Same root cause as the moe_gather_q4/moe_gather_q4_swiglu stubs
  already documented in wh-butter-ops.
- wh-butter-cuda/tests/all_models.rs: dropped a `smallmodel_fuse_slicecast_ab`
  test calling `wh_butter_modeltests::bench_smallmodel_fuse_slicecast` —
  confirmed that function doesn't exist anywhere in wh-butter-modeltests
  on the original tom/feat/cuda-hip-vulkan-backends branch either
  (pre-existing dead reference predating this cherry-pick, not something
  it introduced).
- wh-butter-cuda/tests/f16norm_f32in.rs: dropped — same pre-existing-bug
  class, references a `add_rms_norm_f16norm` op that was never
  implemented on the source branch.
- rust/Cargo.lock: regenerated via `cargo build` (not hand-edited).

* test(rust): env-gate the marlin weight-prep fixture test

BUTTER_MARLIN_FIXTURE_DIR selects the reference-fixture dir; unset gives a
quiet [skipped] instead of unwrapping a hardcoded /tmp path (repo convention:
fixture/model paths come from env vars, absent -> skip).

* fix(rust): adapt strided_col_copy dispatch to the guardless dev kernel

iron@dev's iron_strided_col_copy takes 5 bindings and has no internal idx
guard (the guarded 6-binding variant only exists on an un-landed kernels
feature branch). Drop the extra total binding and cover exactly s*width
threads, using the largest power-of-two block (<=64) dividing the total so
Laguna's non-64-aligned gate-column shapes stay correct.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature New feature or capability

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant